import time
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties

test_path="../data/test.txt"
output_path="result.txt"

def txt2df(train_path):
    print("读取数据集:", train_path)
    data = []
    with open(train_path, 'r') as readfile:
        current_user = None
        for line in readfile:
            line = line.strip()
            if '|' in line:
                current_user = int(line.split('|')[0])
            else:
                item_info = line.split(' ')
                item = int(item_info[0])
                score = int(item_info[2]) if len(item_info) == 3 else np.nan
                data.append([current_user, item, score])
    return pd.DataFrame(data, columns=['user', 'item', 'score'])


train_dataframe = txt2df("../data/train.txt")
test_dataframe = txt2df("../data/test.txt")

alpha = 0.001
lmbda = 0.1
max_iter = 15
F = 500
U = max(train_dataframe['user'].max(), test_dataframe['user'].max()) + 1
I = max(train_dataframe['item'].max(), test_dataframe['item'].max()) + 1
initial_value = np.sqrt(np.mean(train_dataframe['score']) / F)
P = np.full((U, F), initial_value)
Q = np.full((F, I), initial_value)

errors = []
#梯度下降
for iter in range(max_iter):
    start = time.time()
    total_error = 0
    count=0
    for index, line in enumerate(train_dataframe.values):
        count+=1
        user, item, rating = int(line[0]), int(line[1]), line[2]
        eui = rating - np.dot(P[user, :], Q[:, item])
        total_error+=eui**2
      
        Q[:, item] += alpha * (eui * P[user, :] - lmbda * Q[:, item])
        P[user, :] += alpha * (eui * Q[:, item] - lmbda * P[user, :])
    errors.append(np.sqrt(total_error /count))
    end = time.time()
    print(f"迭代轮数 {iter + 1}, RMSE: {errors[-1]:.4f}, 运行时间: {round(end - start, 2)} s")
    
plt.rcParams['font.sans-serif'] = ['SimHei']  
plt.rcParams['axes.unicode_minus'] = False  
# 绘制误差变化图
plt.figure(figsize=(10, 6))
plt.plot(range(1, max_iter + 1), errors, marker='o')
plt.title('梯度下降误差')
plt.xlabel('迭代轮数')
plt.ylabel('RMSE')
plt.grid(True)
plt.show()

with open(test_path, 'r') as testfile, open(output_path, 'w') as outputfile:
    current_user = 0
    for line in testfile:
        if '|' in line:
            outputfile.write(line)
            current_user = int(line.split('|')[0])
        else:
            item_info = line.strip()
            item_info = item_info.split(' ')
            item = int(item_info[0])
            row_res = np.dot(P[current_user, :], Q[:, item])
            res = max(0, min(100, int(row_res)))  
            outputfile.write(f"{item} {res}\n")